#!/usr/bin/env python3

### Please install huaweicloudkms firstly

from huaweicloudsdkcore.auth.credentials import BasicCredentials
from huaweicloudsdkcore.exceptions import exceptions
from huaweicloudsdkcore.http.http_config import HttpConfig
from huaweicloudsdkkms.v2 import KmsClient, EncryptDataRequestBody, EncryptDataRequest, DecryptDataRequestBody, DecryptDataRequest
import argparse

def kms_Client_Init(ak, sk, endpoint, project_id):
    config = HttpConfig.get_default_config()
    config.ignore_ssl_verification = True
    credentials = BasicCredentials(ak, sk, project_id)
    kms_client = KmsClient.new_builder() \
            .with_http_config(config) \
            .with_credentials(credentials) \
            .with_endpoint(endpoint) \
            .build()
    return kms_client

def encrypt_data_test(args):
    print("***************************************Encrypt_Data start*****************************************")
    kms_client = kms_Client_Init(args.ak, args.sk, args.endpoint, args.project_id)
    encrypt_Data_Request_Body = EncryptDataRequestBody(key_id = args.key_id, plain_text = args.plaintext)
    encrypt_Data_Request = EncryptDataRequest(body =  encrypt_Data_Request_Body)
    try:
        res = kms_client.encrypt_data(encrypt_Data_Request)
        print("********************************************Ciphertext********************************************")
        print(res.cipher_text)
    except Exception as e:
        print("Exception : ", e)

def decrypt_data_test(args):
    print("***************************************Dencrypt_Data start****************************************")
    kms_client = kms_Client_Init(args.ak, args.sk, args.endpoint, args.project_id)
    decrypt_Data_Request_Body = DecryptDataRequestBody(cipher_text = args.ciphertext)
    decrypt_Data_Request = DecryptDataRequest(body = decrypt_Data_Request_Body)
    try:
        res = kms_client.decrypt_data(decrypt_Data_Request)
        print("********************************************Plaintext********************************************")
        print(res.plain_text)
    except Exception as e :
        print("Exception :", e)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description = "En/Decrypt data")
    subparser = parser.add_subparsers(description = "En/Decrypt")

    parserEncrypt = subparser.add_parser("encrypt", description = 'encrypt command',
                                        help = "input parameters: ak, sk, endpoint, project_id, plaintext, key_id")
    parserEncrypt.add_argument("-a", "--ak", type = str, help = "Please input AK")
    parserEncrypt.add_argument("-s", "--sk", type = str, help = "Please input SK")
    parserEncrypt.add_argument("-e", "--endpoint", type = str, help = "Please input endpoint")
    parserEncrypt.add_argument("-p", "--plaintext", type = str, help = "Please input plaintext")
    parserEncrypt.add_argument("-w", "--project_id", type = str, help = "Please input project_id")
    parserEncrypt.add_argument("-k", "--key_id", type = str, help = "Please input key_id")
    parserEncrypt.set_defaults(func = encrypt_data_test)

    parserDecrypt = subparser.add_parser("decrypt", description = 'decryption command',
                                        help = "input parameters: ak, sk, endpoint, project_id, ciphertext")
    parserDecrypt.add_argument("-a", "--ak", type = str, help = "Please input AK")
    parserDecrypt.add_argument("-s", "--sk", type = str, help = "Please input SK")
    parserDecrypt.add_argument("-e", "--endpoint", type = str, help = "Please input endpoint")
    parserDecrypt.add_argument("-w", "--project_id", type = str, help = "Please input project_id")
    parserDecrypt.add_argument("-c", "--ciphertext", type = str, help = "Please input ciphertext")
    parserDecrypt.set_defaults(func = decrypt_data_test)


    args = parser.parse_args()
    args.func(args)


